package com.oscar.springaidemo.web.controller;

import com.alibaba.cloud.ai.graph.*;
import com.alibaba.cloud.ai.graph.exception.GraphRunnerException;
import com.alibaba.cloud.ai.graph.exception.GraphStateException;
import com.alibaba.cloud.ai.graph.observation.GraphObservationLifecycleListener;
import io.micrometer.observation.ObservationRegistry;
import io.micrometer.observation.Observations;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.util.HashMap;
import java.util.Map;
import java.util.Optional;

@RestController
@RequestMapping("/graph")
@Slf4j
public class SimpleGraphController {

    private final CompiledGraph compiledGraph;

    public SimpleGraphController(StateGraph stateGraph) throws GraphStateException {
        CompileConfig build = CompileConfig.builder()
                .build();
        this.compiledGraph = stateGraph.compile(build);
    }

    @GetMapping("/expand")
    public Map<String,Object> expand(@RequestParam(value = "query",defaultValue = "介绍下你自己") String query,
                                     @RequestParam(value = "expander_number",defaultValue = "3")Integer number,
                                     @RequestParam(value = "thread_id",defaultValue = "pscar")String threadId) throws GraphRunnerException {
        RunnableConfig build = RunnableConfig.builder().threadId(threadId).build();
        HashMap<String, Object> map = new HashMap<>();
        map.put("query",query);
        map.put("expander_number",number);
        Optional<OverAllState> invoke = this.compiledGraph.invoke(map, build);
        Map<String, Object> stringObjectMap = invoke.map(OverAllState::data).orElse(new HashMap<>());
        return stringObjectMap;


    }

}
